{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2024-12-01T13:31:54.977413Z",
     "start_time": "2024-12-01T13:31:46.925679Z"
    }
   },
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "from torchvision import transforms\n",
    "import numpy as np"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:32:58.917384Z",
     "start_time": "2024-12-01T13:32:58.762135Z"
    }
   },
   "cell_type": "code",
   "source": [
    "train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=transforms.ToTensor())\n",
    "test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=transforms.ToTensor())"
   ],
   "id": "2bdc6c96de67372",
   "outputs": [],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:32:59.188920Z",
     "start_time": "2024-12-01T13:32:59.181760Z"
    }
   },
   "cell_type": "code",
   "source": [
    "train_loader=torch.utils.data.DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)\n",
    "test_loader=torch.utils.data.DataLoader(dataset=test_dataset, batch_size=64, shuffle=True)"
   ],
   "id": "e5dd20259eaa4983",
   "outputs": [],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:32:59.891659Z",
     "start_time": "2024-12-01T13:32:59.807658Z"
    }
   },
   "cell_type": "code",
   "source": [
    "for x, y in train_loader:\n",
    "    #print(y)\n",
    "    print(x.view(x.shape[0], -1).shape)\n",
    "    #print(np.array(x.view(x.shape[0],-1)[0]).shape)\n",
    "    \n",
    "    break"
   ],
   "id": "d19f7af218d98b49",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([64, 784])\n"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:33:00.461875Z",
     "start_time": "2024-12-01T13:33:00.453337Z"
    }
   },
   "cell_type": "code",
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self,hidden_size=128,n_classes=10):\n",
    "        super(Net, self).__init__()\n",
    "        self.fc1 = nn.Linear(in_features=28*28, out_features=hidden_size)\n",
    "        self.fc2 = nn.Linear(in_features=hidden_size, out_features=n_classes)\n",
    "        self.softmax = nn.LogSoftmax(dim=1)\n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = self.relu(self.fc1(x))\n",
    "        x = self.relu(self.fc2(x))\n",
    "        x = self.softmax(x)\n",
    "        return x"
   ],
   "id": "a7b04657901767",
   "outputs": [],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:33:01.286513Z",
     "start_time": "2024-12-01T13:33:01.278889Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model = Net()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)"
   ],
   "id": "21de78cd1ff5a81e",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:33:01.894366Z",
     "start_time": "2024-12-01T13:33:01.885049Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def train(model, train_loader, optimizer,epochs):\n",
    "    model.train()\n",
    "    loss_mem=0\n",
    "    for epoch in range(epochs):\n",
    "        for x, y in train_loader:\n",
    "            x=x.view(x.shape[0], -1) \n",
    "            output = model(x)\n",
    "            loss = criterion(output, y)\n",
    "            if(abs(loss-loss_mem) < 1e-3): #大概就这意思，多来几轮就差不多可以触发了。但正如我所说的，时间来不及了:)\n",
    "                print('stop: in epoch {}'.format(epoch))\n",
    "                optimizer.zero_grad()\n",
    "                break\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        print('Epoch: {}, Loss: {}'.format(epoch, loss.item()))"
   ],
   "id": "6dd657cd8ce63a3f",
   "outputs": [],
   "execution_count": 7
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-01T13:35:09.840878Z",
     "start_time": "2024-12-01T13:33:02.627898Z"
    }
   },
   "cell_type": "code",
   "source": "train(model,train_loader,optimizer,10)",
   "id": "f390f6a6ebeeb511",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss: 0.5131379961967468\n",
      "Epoch: 1, Loss: 0.6989604830741882\n",
      "Epoch: 2, Loss: 0.5990148782730103\n",
      "Epoch: 3, Loss: 0.601474404335022\n",
      "Epoch: 4, Loss: 0.5292885899543762\n",
      "Epoch: 5, Loss: 0.603938639163971\n",
      "Epoch: 6, Loss: 0.5338292717933655\n",
      "Epoch: 7, Loss: 0.7738308906555176\n",
      "Epoch: 8, Loss: 0.4848354756832123\n",
      "Epoch: 9, Loss: 0.525934636592865\n"
     ]
    }
   ],
   "execution_count": 8
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "",
   "id": "ab01271e054b3e6f"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
